package com.zengjianjun.es.demo.highlevelclient;

import com.alibaba.fastjson2.JSON;
import com.zengjianjun.es.demo.config.EsConfig;
import com.zengjianjun.es.demo.entity.EsBaseEntity;
import com.zengjianjun.es.demo.highlevelclient.exception.EsException;
import com.zengjianjun.es.demo.highlevelclient.param.EsQueryParam;
import com.zengjianjun.es.demo.highlevelclient.param.EsUpdateByQueryParam;
import com.zengjianjun.es.demo.highlevelclient.param.ParamQueryBuilder;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.aggregations.AggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.xcontent.XContentType;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * @author zengjianjun
 */
@Slf4j
@Component
public class EsClient {

    @Resource
    private RestHighLevelClient client;

    /**
     * 默认的查询方式
     */
    public SearchResponse defaultQuery(SearchRequest searchRequest) {
        try {
            return client.search(searchRequest, EsConfig.COMMON_OPTIONS);
        } catch (IOException e) {
            throw new EsException("ES defaultQuery 查询失败", e);
        }
    }

    /**
     * 多条件查询列表
     */
    public <T> List<T> queryList(String index, EsQueryParam esQueryParam, Integer from, Integer size, Class<T> resultClazz) {
        try {
            SearchRequest searchRequest = new SearchRequest(index);
            BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
            ParamQueryBuilder.paramBoolQueryBuilder(queryBuilder, esQueryParam);
            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
            sourceBuilder.query(queryBuilder);
            sourceBuilder.from((from - 1) < 0 ? 0 : (from - 1) * size);
            sourceBuilder.size(size);
            searchRequest.source(sourceBuilder);
            SearchResponse searchResponse = client.search(searchRequest, EsConfig.COMMON_OPTIONS);
            RestStatus status = searchResponse.status();
            if (status != RestStatus.OK) {
                return Collections.emptyList();
            }
            List<T> resultList = new ArrayList<>();
            for (SearchHit hit : searchResponse.getHits()) {
                resultList.add(JSON.parseObject(hit.getSourceAsString(), resultClazz));
            }
            return resultList;
        } catch (Exception e) {
            throw new EsException("ES queryList 查询失败", e);
        }
    }

    /**
     * 多条件查询数量
     */
    public Long queryCount(String index, EsQueryParam esQueryParam) {
        try {
            CountRequest countRequest = new CountRequest(index);
            BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
            ParamQueryBuilder.paramBoolQueryBuilder(queryBuilder, esQueryParam);
            SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
            sourceBuilder.query(queryBuilder);
            countRequest.source(sourceBuilder);
            CountResponse countResponse = client.count(countRequest, EsConfig.COMMON_OPTIONS);
            RestStatus status = countResponse.status();
            if (status != RestStatus.OK) {
                return 0L;
            }
            return countResponse.getCount();
        } catch (Exception e) {
            throw new EsException("ES queryCount 查询失败", e);
        }
    }

    /**
     * 常规组合查询
     * <li>size用于指定返回列表中的数据数量，一般情况下只关心聚合结果的话，建议将size指定为0</li>
     */
    public SearchResponse aggregationByBoolQuery(String index, BoolQueryBuilder queryBuilder, AggregationBuilder aggregation, int size) {
        SearchRequest searchRequest = new SearchRequest(index);
        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
        sourceBuilder.query(queryBuilder).aggregation(aggregation);
        sourceBuilder.size(size);
        searchRequest.source(sourceBuilder);
        try {
            return client.search(searchRequest, EsConfig.COMMON_OPTIONS);
        } catch (IOException e) {
            throw new EsException("ES aggregationByBoolQuery 查询失败", e);
        }
    }

    /**
     * 根据条件更新
     */
    public void updateByQueryAsync(EsUpdateByQueryParam param, int slices) {
        UpdateByQueryRequest request = new UpdateByQueryRequest(param.getIndex());
        // 设置分片并行
        request.setSlices(slices);
        // 设置版本冲突时继续执行
        request.setConflicts("proceed");
        // 查询条件如果是and关系使用must，如果是or关系使用should
        BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
        ParamQueryBuilder.paramBoolQueryBuilder(queryBuilder, param.getQueryParam());
        QueryBuilder sourceBuilder = new SearchSourceBuilder().query(queryBuilder).fetchSource(param.getIncludes() != null ? param.getIncludes() : new String[]{}, new String[]{}).query();
        request.setQuery(sourceBuilder);
        request.setScript(new Script(ScriptType.INLINE, "painless", param.getScript(), param.getUpdateFieldAndVal()));
        // 更新最大文档数
//        request.setMaxDocs(10000000);
        // 批次大小
        request.setBatchSize(slices * 1000);
        // 使用滚动参数来控制“搜索上下文”存活的时间
        request.setScroll(TimeValue.timeValueMinutes(10));
        // 超时
        request.setTimeout(TimeValue.timeValueMinutes(20));
        // 设置更新完成后刷新索引 ps：很重要如果不加可能数据不会实时刷新
        request.setRefresh(true);
        client.updateByQueryAsync(request, EsConfig.COMMON_OPTIONS, param.getListener());
    }

    /**
     * 新增
     *
     * @param entity 新增信息
     */
    public <T extends EsBaseEntity> void save(String index, T entity) {
        IndexRequest request = new IndexRequest(index)
                .id(entity.getId())
                .source(JSON.toJSONString(entity), XContentType.JSON);
        BulkRequest bulkRequest = new BulkRequest();
        bulkRequest.add(request);
        this.bulk(bulkRequest);
    }

    /**
     * 批量新增
     *
     * @param entityList 新增信息
     */
    public <T extends EsBaseEntity> void save(String index, List<T> entityList) {
        BulkRequest bulkRequest = new BulkRequest();
        entityList.forEach(entity -> {
            IndexRequest request = new IndexRequest(index)
                    .id(entity.getId())
                    .source(JSON.toJSONString(entity), XContentType.JSON);
            bulkRequest.add(request);
        });
        this.bulk(bulkRequest);
    }

    /**
     * 批量新增
     *
     * @param bulkRequest 新增信息
     */
    public void bulk(BulkRequest bulkRequest) {
        try {
            client.bulk(bulkRequest, EsConfig.COMMON_OPTIONS);
        } catch (Exception e) {
            throw new EsException(e.getMessage());
        }
    }

    /**
     * 异步新增
     *
     * @param index      索引
     * @param entityList 新增信息
     * @param listener   回调监听
     */
    public <T extends EsBaseEntity> void saveAsync(String index, List<T> entityList, ActionListener<BulkResponse> listener) {
        BulkRequest bulkRequest = new BulkRequest();
        entityList.forEach(entity -> {
            IndexRequest request = new IndexRequest(index)
                    .id(entity.getId())
                    .source(JSON.toJSONString(entity), XContentType.JSON);
            bulkRequest.add(request);
        });
        this.bulkAsync(bulkRequest, listener);
    }

    /**
     * 异步批量新增
     *
     * @param bulkRequest 新增信息
     * @param listener    回调监听
     */
    public void bulkAsync(BulkRequest bulkRequest, ActionListener<BulkResponse> listener) {
        client.bulkAsync(bulkRequest, EsConfig.COMMON_OPTIONS, listener);
    }
}
